use std::collections::{HashMap, HashSet};
use std::path::Path;

use anyhow::{Context, anyhow};
use polars::frame::DataFrame;
use polars::functions::concat_df_horizontal;
#[allow(unused_imports)]
use polars::prelude::col;
use polars::prelude::{
    DataFrameOps, IntoLazy, LazyCsvReader, LazyFileListReader, LazyFrame, NullValues, PlPath,
    ToDummies, UnionArgs, concat_lf_horizontal, nth,
};

use crate::parser::Parser;

use super::errors::PlundererError;
use super::parser::{Lexer, Token};
#[derive(Debug, PartialEq)]
pub(crate) struct FileSpec {
    pub(super) file_path: String,
    pub(super) has_header: bool,
    pub(super) fid_spec: FidSpec,
    pub(super) iid_spec: IidSpec,
    pub(super) selections: SelectionSpec,
    pub(super) exclusions: Vec<ColumnSpec>,
}

#[derive(Debug, Clone, PartialEq)]
pub(crate) enum FidSpec {
    Index(usize),
    Str(String),
    Missing,
}

#[derive(Debug, Clone, PartialEq)]
pub(crate) enum IidSpec {
    Index(usize),
    Str(String),
}

#[derive(Debug, Clone, PartialEq)]
pub(crate) enum SelectionSpec {
    All { to_dummy: bool },
    Cols(Vec<SelectCol>),
}

#[derive(Debug, Clone, PartialEq)]
pub(crate) struct SelectCol {
    pub(super) spec: ColumnSpec,
    pub(super) to_dummy: bool,
    pub(super) alias: Option<String>,
}

#[derive(Debug, Clone, PartialEq)]
pub(crate) enum ColumnSpec {
    Index(usize),
    Str(String),
}

impl FileSpec {
    /// Generate AST from tokens.
    #[inline]
    fn from_tokens(tokens: &[Token]) -> Result<Self, PlundererError> {
        Parser::new(tokens).parse_file_spec()
    }

    /// Generate AST from string.
    pub(super) fn from_str<T: AsRef<str>>(s: T) -> Result<Self, PlundererError> {
        let tokens = Lexer::new(s.as_ref()).tokenize();
        Self::from_tokens(&tokens)
    }

    /// Read the file and return the corresponding LazyFrame.
    pub(super) fn lazy_read_to_lazyframe(&self) -> Result<LazyFrame, anyhow::Error> {
        // check file existence
        let path = Path::new(&self.file_path);
        if !path.exists() {
            return Err(PlundererError::FileNotFoundError(self.file_path.clone()).into());
        }

        if !path.exists() {
            return Err(PlundererError::FileNotFoundError(self.file_path.clone()).into());
        }

        // // Read the first line of the file to get the header
        // let mut file_reader = File::open(path)
        //     .with_context(|| anyhow!("Failed to create a reader for {}", path.display()))?;

        // let source_headers:  = {
        //     let mut first_line = String::new();
        //     std::io::BufReader::new(&mut file_reader)
        //         .read_line(&mut first_line)
        //         .with_context(|| anyhow!("Failed to read first line of {}", path.display()))?;
        //     first_line
        // };

        // Establish a LazyFrame
        log::debug!("Creating lazy csv reader...");
        let mut source_lf = LazyCsvReader::new(PlPath::from_str(&self.file_path))
            .with_has_header(self.has_header)
            .with_separator(b'\t')
            .with_rechunk(true)
            .with_null_values(Some(NullValues::AllColumns(vec![
                "NAN".into(),
                "NaN".into(),
                "Nan".into(),
                "nan".into(),
                "NA".into(),
                "Na".into(),
                "na".into(),
                "NULL".into(),
                "Null".into(),
                "null".into(),
            ])))
            .with_infer_schema_length(None) // Though slow, but necessary for accurate schema inference
            .finish()?;

        log::debug!("Getting schema...");
        let source_schema = source_lf.collect_schema()?;
        let mut source_headers_map: HashMap<String, usize> = HashMap::new();

        for (index, header_name) in source_schema.iter_names().enumerate() {
            source_headers_map.insert(header_name.to_string(), index);
        }
        log::trace!("columns are: {:?}", source_headers_map);
        log::debug!("Schema got.");
        let columns_count = source_schema.len();
        assert_eq!(columns_count, source_headers_map.len());

        let fid_col_idx: Option<usize>;
        let iid_col_idx: usize;

        match self.fid_spec {
            FidSpec::Missing => {
                fid_col_idx = None;
            }
            FidSpec::Index(idx) => {
                fid_col_idx = Some(idx);
            }
            FidSpec::Str(ref s) => {
                if let Some(idx) = source_headers_map.get(s) {
                    fid_col_idx = Some(*idx);
                } else {
                    return Err(PlundererError::ColumnNotFoundError(s.to_string()).into());
                }
            }
        }
        match self.iid_spec {
            IidSpec::Index(idx) => {
                iid_col_idx = idx;
            }
            IidSpec::Str(ref s) => {
                if let Some(idx) = source_headers_map.get(s) {
                    iid_col_idx = *idx;
                } else {
                    return Err(PlundererError::ColumnNotFoundError(s.to_string()).into());
                }
            }
        }

        let mut exclusions_set: HashSet<usize> = HashSet::new();
        for exclusion in &self.exclusions {
            match exclusion {
                ColumnSpec::Index(idx) => {
                    exclusions_set.insert(*idx);
                }
                ColumnSpec::Str(s) => {
                    if !self.has_header {
                        return Err(PlundererError::ColumnHeaderDesignationError.into());
                    }
                    if let Some(idx) = source_headers_map.get(s) {
                        exclusions_set.insert(*idx);
                    } else {
                        return Err(PlundererError::ColumnNotFoundError(s.to_string()).into());
                    }
                }
            }
        }

        let covar_only_lf: LazyFrame = match self.selections {
            SelectionSpec::All { to_dummy } => {
                log::debug!("Select all columns of {}", path.display());
                let mut lf = source_lf.clone().select({
                    let mut exprs = Vec::new();
                    for i in 0..source_headers_map.len() {
                        let is_fid = match fid_col_idx {
                            Some(idx) => {
                                if i == idx {
                                    true
                                } else {
                                    false
                                }
                            }
                            None => false,
                        };
                        let is_iid = { i == iid_col_idx };
                        if !is_fid && !is_iid && !exclusions_set.contains(&i) {
                            exprs.push(nth(i as i64).as_expr())
                        }
                    }
                    exprs
                });
                if to_dummy {
                    log::debug!("Converting all columns to dummy variables");
                    lf = lf
                        .collect()
                        .with_context(|| anyhow!("Failed to collect data"))?
                        .to_dummies(Some("_"), true, false)
                        .with_context(|| {
                            anyhow!("Failed to convert all columns to dummy variables.")
                        })?
                        .lazy();
                };
                lf
            }
            SelectionSpec::Cols(ref select_col_vec) => {
                let mut covar_dfs: Vec<LazyFrame> = Vec::new();
                for select_col in select_col_vec {
                    log::debug!("Select column {:?}", select_col);
                    match &select_col.spec {
                        ColumnSpec::Index(idx) => {
                            log::debug!("Select column index {idx}");
                            if exclusions_set.contains(&idx) {
                                log::debug!("Exclusion set contains {idx}. It will be excluded.");
                                continue;
                            } else {
                                let exprs = match select_col.alias {
                                    Some(ref alias) => [nth(*idx as i64).as_expr().alias(alias)],
                                    None => [nth(*idx as i64).as_expr()],
                                };
                                let this_col_lf = source_lf.clone().select(exprs);
                                if select_col.to_dummy {
                                    log::debug!("Converting column index {idx} to dummy...");
                                    covar_dfs.push(
                                        this_col_lf
                                            .collect()
                                            .with_context(|| {
                                                anyhow!("Failed to collect lazyframe for index {idx}")
                                            })?
                                            .to_dummies(Some("_"), true, false)
                                            .with_context(|| anyhow!("Failed to convert all covariant columns to dummy variables!"))?
                                            .lazy(),
                                    );
                                } else {
                                    covar_dfs.push(this_col_lf);
                                }
                            }
                        }
                        ColumnSpec::Str(s) => {
                            if !self.has_header {
                                return Err(PlundererError::ColumnHeaderDesignationError.into());
                            }
                            let idx = source_headers_map
                                .get(s)
                                .with_context(|| {
                                    anyhow!("Failed to get {s} from {source_headers_map:?}")
                                })?
                                .clone();
                            log::debug!("Select column idx {idx}, header {s}");
                            if exclusions_set.contains(&idx) {
                                log::debug!(
                                    "Exclusion set contains idx {idx}. It will be excluded."
                                );
                                continue;
                            } else {
                                let exprs = match select_col.alias {
                                    Some(ref alias) => [nth(idx as i64).as_expr().alias(alias)],
                                    None => {
                                        if !self.has_header {
                                            [nth(idx as i64).as_expr().alias(format!(
                                                "{}_{}",
                                                path.file_name().unwrap().to_str().unwrap(),
                                                idx
                                            ))]
                                        } else {
                                            [nth(idx as i64).as_expr()]
                                        }
                                    }
                                };
                                let this_col_lf = source_lf.clone().select(exprs);
                                if select_col.to_dummy {
                                    log::debug!("Converting column {s}, index {idx} to dummy");
                                    covar_dfs.push(
                                        this_col_lf
                                            .collect()
                                            .with_context(|| {
                                                anyhow!("Failed to collect data for column {s}, index {idx}")
                                            })?
                                            .to_dummies(Some("_"), true, false)
                                            .with_context(|| {
                                                anyhow!("Failed to convert column {s} to dummy")
                                            })?
                                            .lazy(),
                                    );
                                } else {
                                    covar_dfs.push(this_col_lf);
                                }
                            }
                        }
                    }
                }
                log::debug!("Concatenating covariate dataframes...");
                concat_lf_horizontal(covar_dfs, UnionArgs::default())?
            }
        };

        log::debug!("Select FID column");
        let fid_lf = match fid_col_idx {
            Some(idx) => source_lf
                .clone()
                .select([nth(idx as i64).as_expr().alias("FID")]),
            None => source_lf
                .clone()
                .select([nth(iid_col_idx as i64).as_expr().alias("FID")]),
        };
        log::debug!("Select IID column");
        let iid_lf = source_lf
            .clone()
            .select([nth(iid_col_idx as i64).as_expr().alias("IID")]);

        log::debug!("Concatenating FID, IID, and covariate dataframes...");
        let res = concat_lf_horizontal(&[fid_lf, iid_lf, covar_only_lf], UnionArgs::default())?;

        log::debug!("LazyFrame from {} is returned", path.display());
        Ok(res)
    }

    /// This method is deprecated. It doesn't utilize the LazyFrame API. Use
    /// `lazy_read_to_lazyframe` instead.
    #[deprecated(since = "0.1.1", note = "Use `lazy_read_to_lazyframe` instead.")]
    #[warn(deprecated)]
    #[allow(dead_code)]
    pub(super) fn read_to_lazyframe(&self) -> Result<LazyFrame, anyhow::Error> {
        // check file existence
        let path = Path::new(&self.file_path);
        if !path.exists() {
            return Err(PlundererError::FileNotFoundError(self.file_path.clone()).into());
        }

        let lf = LazyCsvReader::new(PlPath::from_str(&self.file_path))
            .with_has_header(self.has_header)
            .with_separator(b'\t')
            .with_rechunk(true)
            .with_null_values(Some(NullValues::AllColumns(vec![
                "NAN".into(),
                "NaN".into(),
                "Nan".into(),
                "nan".into(),
                "NA".into(),
                "Na".into(),
                "na".into(),
                "NULL".into(),
                "Null".into(),
                "null".into(),
            ])))
            .with_infer_schema_length(None)
            .finish()?;

        let source_df = lf.collect()?;

        let source_columns = source_df.get_columns();

        let source_headers: Vec<&str> = source_df
            .get_column_names()
            .iter()
            .map(|s| s.as_str())
            .collect();

        // change source_headers to a HashMap<String, usize>
        let source_headers_map: HashMap<String, usize> = source_headers
            .iter()
            .enumerate()
            .map(|(i, &header)| (header.to_string(), i))
            .collect();

        let fid_col_header: Option<String>;
        let iid_col_header: String;
        let fid_col_idx: Option<usize>;
        let iid_col_idx: usize;
        let mut exclude_indices: HashSet<usize> = HashSet::new();

        match self.fid_spec {
            FidSpec::Missing => {
                fid_col_header = None;
                fid_col_idx = None;
            }
            FidSpec::Index(idx) => {
                fid_col_header = Some(source_headers[idx].to_string());
                fid_col_idx = Some(idx);
            }
            FidSpec::Str(ref s) => {
                if let Some(idx) = source_headers_map.get(s) {
                    fid_col_header = Some(s.to_string());
                    fid_col_idx = Some(*idx);
                } else {
                    return Err(PlundererError::ColumnNotFoundError(s.to_string()).into());
                }
            }
        };
        match self.iid_spec {
            IidSpec::Index(idx) => {
                iid_col_header = source_headers[idx].to_string();
                iid_col_idx = idx;
            }
            IidSpec::Str(ref s) => {
                if let Some(idx) = source_headers_map.get(s) {
                    iid_col_header = s.to_string();
                    iid_col_idx = *idx;
                } else {
                    return Err(PlundererError::ColumnNotFoundError(s.to_string()).into());
                }
            }
        }
        for exclude_col in self.exclusions.iter() {
            match exclude_col {
                ColumnSpec::Index(idx) => {
                    exclude_indices.insert(*idx);
                }
                ColumnSpec::Str(s) => {
                    if let Some(idx) = source_headers_map.get(s) {
                        exclude_indices.insert(*idx);
                    } else {
                        return Err(PlundererError::ColumnNotFoundError(s.to_string()).into());
                    }
                }
            }
        }

        let covar_only_df: DataFrame = match self.selections {
            SelectionSpec::All { to_dummy } => {
                let covar_columns = source_columns
                    .iter()
                    .enumerate()
                    .filter_map(|(idx, col)| {
                        if idx != iid_col_idx
                            && {
                                if let Some(fid_idx) = fid_col_idx {
                                    idx != fid_idx
                                } else {
                                    true
                                }
                            }
                            && !exclude_indices.contains(&idx)
                        {
                            Some(col.clone())
                        } else {
                            None
                        }
                    })
                    .collect::<Vec<_>>();
                let res = DataFrame::new(covar_columns).unwrap();
                if to_dummy {
                    res.to_dummies(Some("_"), true, false).with_context(|| {
                        anyhow!("Failed to convert all columns to dummy variables.")
                    })?
                } else {
                    res
                }
            }
            SelectionSpec::Cols(ref select_col_vec) => {
                let covar_dfs: Vec<DataFrame> = select_col_vec
                    .iter()
                    .filter_map(|select_col| match select_col.spec {
                        ColumnSpec::Index(idx) => {
                            if exclude_indices.contains(&idx) {
                                None
                            } else {
                                let mut this_col = source_columns[idx].clone();
                                let this_df: DataFrame;
                                if let Some(ref alias) = select_col.alias {
                                    this_col.rename(alias.into());
                                }
                                if select_col.to_dummy {
                                    this_df = this_col
                                        .as_materialized_series()
                                        .to_dummies(Some("_"), true, false)
                                        .unwrap();
                                } else {
                                    this_df = this_col.into_frame();
                                }
                                Some(this_df)
                            }
                        }
                        ColumnSpec::Str(ref s) => {
                            let idx = source_headers_map
                                .get(s)
                                .with_context(|| {
                                    anyhow!("Failed to get {s} from {source_headers_map:?}")
                                })
                                .unwrap()
                                .clone();
                            if exclude_indices.contains(&idx) {
                                None
                            } else {
                                let mut this_col = source_columns[idx].clone();
                                let this_df: DataFrame;
                                if let Some(ref alias) = select_col.alias {
                                    this_col.rename(alias.into());
                                }
                                if select_col.to_dummy {
                                    this_df = this_col
                                        .as_materialized_series()
                                        .to_dummies(Some("_"), true, false)
                                        .unwrap();
                                } else {
                                    this_df = this_col.into_frame();
                                }
                                Some(this_df)
                            }
                        }
                    })
                    .collect();
                let covar_df = concat_df_horizontal(&covar_dfs, true)?;
                covar_df
            }
        };

        let mut all_dfs_vec = Vec::new();
        if let Some(fid_idx) = fid_col_idx {
            let mut fid_df = source_columns[fid_idx].clone();
            fid_df.rename("FID".into());
            all_dfs_vec.push(fid_df.into_frame());
        } else {
            // Copy iid column to fid
            let mut fid_col = source_columns[iid_col_idx].clone();
            fid_col.rename("FID".into());
            all_dfs_vec.push(fid_col.into_frame());
        }
        all_dfs_vec.push({
            let mut iid_df = source_columns[iid_col_idx].clone();
            iid_df.rename("IID".into());
            iid_df.into_frame()
        });
        all_dfs_vec.push(covar_only_df);

        let res = concat_df_horizontal(&all_dfs_vec, true)?;

        Ok(res.lazy())
    }
}

#[cfg(test)]
mod file_tests {
    use std::{cell::LazyCell, path::PathBuf};

    use polars::df;

    use super::*;

    const FID_AND_IID_FILE: &str = "testfiles/fid_and_iid_file1.tsv";
    const FID_AND_IID_FILE_DF: LazyCell<DataFrame> = LazyCell::new(|| {
        df!(
            "FID" => ["family_1", "family_1", "family_2", "family_2"],
            "IID" => ["person_1", "person_2", "person_1", "person_3"],
            "Age" => [25, 8, 86, 35],
            "Sex" => ["female", "male", "female", "female"],
            "Ethnicity" => ["white", "mixed", "chinese", "chinese"],
            "Weight" => [75.64, 22.3, 56.2, 53.2],
            "Height" => [168.4, 104., 165., 166.],
            "High Blood Pressure" => [153, 110, 120, 90],
            "Low Blood Pressure" => [115, 74, 89, 60]
        )
        .unwrap()
    });

    #[test]
    fn test_fid_and_iid_file() {
        // Test file function
        assert!(
            PathBuf::from(FID_AND_IID_FILE).exists(),
            "{FID_AND_IID_FILE} does not exists!"
        );

        let file_spec = FileSpec {
            file_path: FID_AND_IID_FILE.to_string(),
            has_header: true,
            fid_spec: FidSpec::Index(0),
            iid_spec: IidSpec::Index(1),
            selections: SelectionSpec::All { to_dummy: false },
            exclusions: Vec::new(),
        };

        let expected_df = &*FID_AND_IID_FILE_DF;

        let res_df = file_spec
            .lazy_read_to_lazyframe()
            .unwrap()
            .collect()
            .unwrap();

        assert_eq!(&res_df, expected_df);
    }

    #[test]
    fn test_fid_and_iid_file_with_exclusions() {
        let file_spec = FileSpec {
            file_path: FID_AND_IID_FILE.to_string(),
            has_header: true,
            fid_spec: FidSpec::Index(0),
            iid_spec: IidSpec::Index(1),
            selections: SelectionSpec::All { to_dummy: false },
            exclusions: vec![
                ColumnSpec::Str("Age".into()), // Index 3
                ColumnSpec::Index(3),          // Sex
            ],
        };

        let expected_df = FID_AND_IID_FILE_DF.drop_many(["Age", "Sex"]);

        let res_df = file_spec
            .lazy_read_to_lazyframe()
            .unwrap()
            .collect()
            .unwrap();

        assert_eq!(res_df, expected_df);
    }

    #[test]
    fn test_fid_and_iid_file_with_selections() {
        let file_spec = FileSpec {
            file_path: FID_AND_IID_FILE.to_string(),
            has_header: true,
            fid_spec: FidSpec::Index(0),
            iid_spec: IidSpec::Index(1),
            selections: SelectionSpec::Cols(vec![
                SelectCol {
                    spec: ColumnSpec::Index(4), // Ethnicity
                    alias: Some(String::from("ETHNICITY")),
                    to_dummy: false,
                },
                SelectCol {
                    spec: ColumnSpec::Str(String::from("Height")), // Index 6
                    alias: Some(String::from("HEIGHT")),
                    to_dummy: false,
                },
            ]),
            exclusions: Vec::new(),
        };

        let expected_df = FID_AND_IID_FILE_DF
            .clone()
            .lazy()
            .select([col("FID"), col("IID"), col("Ethnicity"), col("Height")])
            .collect()
            .unwrap();

        let res_df = file_spec
            .lazy_read_to_lazyframe()
            .unwrap()
            .collect()
            .unwrap();

        assert_eq!(res_df, expected_df);
    }

    #[test]
    fn test_fid_and_iid_file_with_selections_and_exclusions() {
        let file_spec = FileSpec {
            file_path: FID_AND_IID_FILE.to_string(),
            has_header: true,
            fid_spec: FidSpec::Index(0),
            iid_spec: IidSpec::Index(1),
            selections: SelectionSpec::Cols(vec![
                SelectCol {
                    spec: ColumnSpec::Index(4), // Ethnicity
                    alias: Some(String::from("ETHNICITY")),
                    to_dummy: false,
                },
                SelectCol {
                    spec: ColumnSpec::Str(String::from("Height")), // Index 6
                    alias: Some(String::from("HEIGHT")),
                    to_dummy: false,
                },
                SelectCol {
                    spec: ColumnSpec::Str(String::from("Age")), // Index 2
                    alias: None,
                    to_dummy: false,
                },
            ]),
            exclusions: vec![
                ColumnSpec::Index(2),
                ColumnSpec::Str("Ethnicity".to_string()),
            ],
        };

        let expected_df = FID_AND_IID_FILE_DF
            .clone()
            .lazy()
            .select([col("FID"), col("IID"), col("Height")])
            .collect()
            .unwrap();

        let res_df = file_spec
            .lazy_read_to_lazyframe()
            .unwrap()
            .collect()
            .unwrap();

        assert_eq!(res_df, expected_df);
    }

    #[test]
    fn test_fid_and_iid_file_all_to_dummy() {
        let file_spec = FileSpec {
            file_path: FID_AND_IID_FILE.to_string(),
            has_header: true,
            fid_spec: FidSpec::Index(0),
            iid_spec: IidSpec::Index(1),
            selections: SelectionSpec::All { to_dummy: true },
            exclusions: Vec::new(),
        };

        let _res_df = file_spec.read_to_lazyframe().unwrap().collect().unwrap();
    }

    #[test]
    fn test_fid_and_iid_file_with_selections_to_dummy() {
        let file_spec = FileSpec {
            file_path: FID_AND_IID_FILE.to_string(),
            has_header: true,
            fid_spec: FidSpec::Index(0),
            iid_spec: IidSpec::Index(1),
            selections: SelectionSpec::Cols(vec![
                SelectCol {
                    spec: ColumnSpec::Index(4), // Ethnicity
                    alias: None,
                    to_dummy: true,
                },
                SelectCol {
                    spec: ColumnSpec::Str(String::from("Height")), // Index 6
                    alias: None,
                    to_dummy: true,
                },
            ]),
            exclusions: Vec::new(),
        };

        let res_df = file_spec
            .lazy_read_to_lazyframe()
            .unwrap()
            .collect()
            .unwrap();

        let expected_df = concat_df_horizontal(
            &[
                FID_AND_IID_FILE_DF.select(["FID", "IID"]).unwrap(),
                FID_AND_IID_FILE_DF
                    .select(["Ethnicity", "Height"])
                    .unwrap()
                    .to_dummies(Some("_"), true, false)
                    .unwrap(),
            ],
            true,
        )
        .unwrap();

        assert_eq!(expected_df, res_df)
    }
}
